#include "QR.h"
#include <iostream>

using namespace std;
using namespace RSIM;

////////////////////////////////////////////////////////////////////////

double RSIM::hypot(const double &a, const double &b){
	if(a==0)
		return abs(b);
	else{   
		double c = b/a;
		return fabs(a) * sqrt(1 + c*c);
	}   
}

////////////////////////////////////////////////////////////////////////

RSIM::QR::QR(const Matrix &A):QR_(A){
	
	m = A.getNRow();
	n = A.getNCol();
	Rdiag = Vector(n);
	int i=0, j=0, k=0;
	
	// Main loop.
	for (k = 0; k < n; k++) {
		// Compute 2-norm of k-th column without under/overflow.
		double nrm = 0;
		for (i = k; i < m; i++) {
			nrm = hypot(nrm,QR_[i][k]);
		}
		
		if (nrm != 0.0) {
			// Form k-th Householder vector.
			if (QR_[k][k] < 0) {
				nrm = -nrm;
			}
			for (i = k; i < m; i++) {
				QR_[i][k] /= nrm;
			}
			QR_[k][k] += 1.0;
			
			// Apply transformation to remaining columns.
			for (j = k+1; j < n; j++) {
				double s = 0.0; 
				for (i = k; i < m; i++) {
					s += QR_[i][k]*QR_[i][j];
				}
				s = -s/QR_[k][k];
				for (i = k; i < m; i++) {
					QR_[i][j] += s*QR_[i][k];
				}
			}
		}
		Rdiag[k] = -nrm;
	}
}

////////////////////////////////////////////////////////////////////////

int QR::isFullRank() const		
{
	for (int j = 0; j < n; j++) 
	{
		if (Rdiag[j] == 0)
			return 0;
	}
	return 1;
}

////////////////////////////////////////////////////////////////////////

Matrix QR::getHouseholder(void)const{
	Matrix H(m,n);
	
	/* note: H is completely filled in by algorithm, so
	initializaiton of H is not necessary.
	*/
	for (int i = 0; i < m; i++) 
	{
		for (int j = 0; j < n; j++) 
		{
			if (i >= j) {
				H[i][j] = QR_[i][j];
			} else {
				H[i][j] = 0.0;
			}
		}
	}
	return H;
}

////////////////////////////////////////////////////////////////////////

Matrix QR::getR() const
{
	Matrix R(n,n);
	for (int i = 0; i < n; i++) {
		for (int j = 0; j < n; j++) {
			if (i < j) {
				R[i][j] = QR_[i][j];
			} else if (i == j) {
				R[i][j] = Rdiag[i];
			} else {
				R[i][j] = 0.0;
			}
		}
	}
	return R;
}

////////////////////////////////////////////////////////////////////////

Matrix QR::getQ() const
{
	int i=0, j=0, k=0;
	
	Matrix Q(m,n);
	for (k = n-1; k >= 0; k--) {
		for (i = 0; i < m; i++) {
			Q[i][k] = 0.0;
		}
		Q[k][k] = 1.0;
		for (j = k; j < n; j++) {
			if (QR_[k][k] != 0) {
				double s = 0.0;
				for (i = k; i < m; i++) {
					s += QR_[i][k]*Q[i][j];
				}
				s = -s/QR_[k][k];
				for (i = k; i < m; i++) {
					Q[i][j] += s*QR_[i][k];
				}
			}
		}
	}
	return Q;
}

////////////////////////////////////////////////////////////////////////

Vector QR::solve(const Vector &b) const
{
	if (b.getLength() != m)		/* arrays must be conformant */
		return Vector();
	
	if ( !isFullRank() )		/* matrix is rank deficient */
	{
		return Vector();
	}
	
	Vector x = b;
	
	// Compute Y = transpose(Q)*b
	for (int k = 0; k < n; k++) 
	{
		double s = 0.0; 
		for (int i = k; i < m; i++) 
		{
			s += QR_[i][k]*x[i];
		}
		s = -s/QR_[k][k];
		for (int i = k; i < m; i++) 
		{
			x[i] += s*QR_[i][k];
		}
	}
	// Solve R*X = Y;
	for (int k = n-1; k >= 0; k--) 
	{
		x[k] /= Rdiag[k];
		for (int i = 0; i < k; i++) {
			x[i] -= x[k]*QR_[i][k];
		}
	}
	
	
	/* return n x nx portion of X */
	Vector x_(n);
	for (int i=0; i<n; i++)
		x_[i] = x[i];
	
	return x_;
}

////////////////////////////////////////////////////////////////////////

Matrix QR::solve(const Matrix &B) const
{
	if (B.getNRow() != m)		/* arrays must be conformant */
		return Matrix();
	
	if ( !isFullRank() )		/* matrix is rank deficient */
	{
		return Matrix();
	}
	
	int nx = B.getNCol(); 
	Matrix X = B;
	int i=0, j=0, k=0;
	
	// Compute Y = transpose(Q)*B
	for (k = 0; k < n; k++) {
		for (j = 0; j < nx; j++) {
			double s = 0.0; 
			for (i = k; i < m; i++) {
				s += QR_[i][k]*X[i][j];
			}
			s = -s/QR_[k][k];
			for (i = k; i < m; i++) {
				X[i][j] += s*QR_[i][k];
			}
		}
	}
	// Solve R*X = Y;
	for (k = n-1; k >= 0; k--) {
		for (j = 0; j < nx; j++) {
			X[k][j] /= Rdiag[k];
		}
		for (i = 0; i < k; i++) {
			for (j = 0; j < nx; j++) {
				X[i][j] -= X[k][j]*QR_[i][k];
			}
		}
	}
	
	
	/* return n x nx portion of X */
	Matrix X_(n,nx);
	for (i=0; i<n; i++)
		for (j=0; j<nx; j++)
			X_[i][j] = X[i][j];
		
		return X_;
}
